import pycuda.autoinit
import pycuda.driver as cuda
import numpy as np
#计时器
from timeit import default_timer as timer
from pycuda.compiler import SourceModule

from EM_GMM import GMM

class GMM_GPU(GMM):
    def train(self,sample,time):
#载入kernal
        fo = open("gmmGpu.cu", "rb")
        content = str(fo.read(), encoding = "utf8")
        mod = SourceModule(content)
        reduce0 = mod.get_function("reduce0")
        reduce1 = mod.get_function("reduce1")
#
        self.sampleData = sample

        self.sampleData = self.sampleData.astype(np.float32)
        self.weights = self.weights.astype(np.float32)
        self.mu = self.mu.astype(np.float32)
        self.sigma = self.sigma.astype(np.float32)

        grid_size = 1024
        block_size = 1024
 
        sample_gpu = cuda.mem_alloc(self.sampleData.nbytes)
        weights_gpu = cuda.mem_alloc(self.weights.nbytes)
        mu_gpu = cuda.mem_alloc(self.mu.nbytes)
        sigma_gpu = cuda.mem_alloc(self.sigma.nbytes)

        cuda.memcpy_htod(sample_gpu, self.sampleData)

        gamma_gpu_temp = cuda.mem_alloc(grid_size * self.gaussian_num * 4)
        mu_gpu_temp = cuda.mem_alloc(grid_size * self.gaussian_num * 4)
        sigma_gpu_temp = cuda.mem_alloc(grid_size * self.gaussian_num * 4)

        gamma_gpu_out = cuda.mem_alloc(self.gaussian_num * 4)
        mu_gpu_out = cuda.mem_alloc(self.gaussian_num * 4)
        sigma_gpu_out = cuda.mem_alloc(self.gaussian_num * 4)

        gamma_out = np.zeros(self.gaussian_num).astype(np.float32)
        mu_out = np.zeros(self.gaussian_num).astype(np.float32)
        sigma_out = np.zeros(self.gaussian_num).astype(np.float32)

        start = timer()
        N = len(sample)
        for index in range(time):

            cuda.memcpy_htod(weights_gpu, self.weights)
            cuda.memcpy_htod(mu_gpu, self.mu)
            cuda.memcpy_htod(sigma_gpu, self.sigma)
            
            reduce0(mu_gpu, sigma_gpu,weights_gpu,sample_gpu,
            gamma_gpu_temp,mu_gpu_temp,sigma_gpu_temp, block=(block_size,1,1), grid=(grid_size,1))
            reduce1(gamma_gpu_temp,mu_gpu_temp,sigma_gpu_temp,gamma_gpu_out,mu_gpu_out,sigma_gpu_out, block=(block_size,1,1), grid=(1,1))

            cuda.memcpy_dtoh(gamma_out, gamma_gpu_out)
            cuda.memcpy_dtoh(mu_out, mu_gpu_out)
            cuda.memcpy_dtoh(sigma_out, sigma_gpu_out)
            
            self.mu = (mu_out / gamma_out).astype(np.float32)
            self.sigma = np.clip(sigma_out / gamma_out,0.1,3).astype(np.float32)
            self.weights = (gamma_out / N).astype(np.float32)

            # print(self.mu)
            # print(self.sigma)
            # print(self.weights)

            # cuda.memcpy_dtoh(self.sampleData, sample_gpu)
            # print(self.sampleData[0:10])
            #mu_gpu_temp =
        run_time = timer() - start
        print("gpu run time %f seconds " % (run_time))


if __name__ == '__main__':
    #目标
    gmm = GMM(10)
    #gmm.show()
    samples = gmm.sample(1024*1024)
    gmm2 = GMM_GPU(20)
    gmm2.train(samples,1)
    gmm2.show(gmm)

    gmm3 = GMM(20)
    gmm3.train(samples,1)
    gmm3.show(gmm)